import cv2
import numpy as np

img = cv2.imread("./img/color_recognize3.png")
h, w, _ = img.shape
# 设置一个红色区域掩膜
red_mask = np.zeros((h, w)).astype(np.uint8)
# 从BGR颜色标准，转化为HSV颜色标准
img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

# 设置掩膜
for i in range(h):
    for j in range(w):
        h = img_hsv[i, j, 0]
        s = img_hsv[i, j, 1]
        v = img_hsv[i, j, 2]
        # 判断红色
        if (0 <= h <= 10 or 156 <= h <= 180) and \
                43 <= s <= 255 and \
                46 <= v <= 255:
            red_mask[i, j] = 255

# -----------------------将原本的显示红色的区域进行形态学变化-------
kernel = np.ones((5, 5))
# 将mask（掩膜）进行白色像素腐蚀： 除去原图像中离散的颜色像素点
red_mask = cv2.erode(red_mask, kernel)
# 将mask（掩膜）进行白色像素膨胀： 将离散的部分尽可能连接起来
red_mask = cv2.dilate(red_mask, kernel)

# 寻找黑白图中的轮廓点
contours, _ = cv2.findContours(red_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# 找出近似矩形
for contour in contours:
    # 获取近似矩形的参数
    x, y, w, h = cv2.boundingRect(contour)
    # 在原图中绘制矩形以及文字标注
    cv2.rectangle(img, (x, y), (x + w, y + h), color=(0, 0, 0), thickness=2)
    cv2.rectangle(img, (x, y - 20), (x + 50, y), color=(0, 0, 0), thickness=cv2.FILLED)
    cv2.putText(img, "red", (x, y), 0, 1, color=(255, 255, 255), thickness=2)

cv2.imshow("img", img)
# cv2.imshow("mask", red_mask)
cv2.waitKey(0)
